#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/5/8 11:03
# @File    : windos.py
# @Description : -i http://pypi.tuna.tsinghua.edu.cn/simple/ --trusted-host pypi.tuna.tsinghua.edu.cn


import time
import win32api
import win32gui
import win32con
import win32ui

import time
import cv2
import numpy as np
import matplotlib.pyplot as plt
from pykeyboard import ReplyKeyboard
from vk_code import vk_to_char
from ctypes import windll
import matplotlib
matplotlib.use('TkAgg')

class KM_windows:

    def __init__(self, hwnd):
        self.hwnd = hwnd
        self.setKeypadDelay()
        self.setMouseDelay()
        self.x = 0
        self.y = 0

    # 键盘 ===========================
    def setKeypadDelay(self, delay=None):
        self.keyboard_delay = 0.01 if delay is None else delay / 1000

    def keyDownChar(self, key_str):
        win32api.SendMessage(self.hwnd, win32con.WM_KEYDOWN, vk_to_char(key_str), 0)

    def keyUpChar(self, key_str):
        win32api.SendMessage(self.hwnd, win32con.WM_KEYUP, vk_to_char(key_str), 0xC01E0001)

    def keyPressChar(self, key_str: str):
        self.keyDownChar(key_str)
        time.sleep(self.keyboard_delay)
        self.keyUpChar(key_str)

    # def enableRealKeypad(self, enable):
    #     pass

    def sendString(self, hwnd, str_):
        for ch in str_:
            win32api.SendMessage(hwnd, win32con.WM_CHAR, ord(ch), 0)
            time.sleep(self.keyboard_delay)

    # 鼠标 ===================================
    def leftDown(self):
        win32api.SendMessage(self.hwnd, win32con.WM_LBUTTONDOWN, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))  # 鼠标按下

    def leftUp(self):
        win32api.SendMessage(self.hwnd, win32con.WM_LBUTTONUP, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))  # 鼠标按下

    def leftClick(self):
        self.leftDown()
        time.sleep(self.mouse_delay)
        self.leftUp()

    def leftDoubleClick(self):
        self.leftClick()
        time.sleep(self.mouse_delay)
        self.leftClick()

    def rightDown(self):
        win32api.SendMessage(self.hwnd, win32con.WM_RBUTTONDOWN, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))

    def rightUp(self):
        win32api.SendMessage(self.hwnd, win32con.WM_RBUTTONUP, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))

    def rightClick(self):
        self.rightDown()
        time.sleep(self.mouse_delay)
        self.rightUp()

    def setMouseDelay(self, delay=None):
        self.mouse_delay = 0.03 if delay is None else delay / 1000

    def middleDown(self):
        win32api.SendMessage(self.hwnd, win32con.WM_MBUTTONDOWN, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))

    def middleUp(self):
        win32api.SendMessage(self.hwnd, win32con.WM_MBUTTONUP, win32con.MK_LBUTTON,
                             win32api.MAKELONG(self.x, self.y))

    def moveTo(self, x, y):
        self.x = x
        self.y = y

    def wheelDown(self):
        win32api.SendMessage(self.hwnd, win32con.WM_MOUSEWHEEL, win32con.WHEEL_DELTA * -1,
                             win32api.MAKELONG(self.x, self.y))

    def wheelUp(self):
        win32api.SendMessage(self.hwnd, win32con.WM_MOUSEWHEEL, win32con.WHEEL_DELTA * 1,
                             win32api.MAKELONG(self.x, self.y))

    def hotKey(self, key_list, interval=None):
        if not interval:
            interval = self.keyboard_delay
        for key in key_list:
            self.keyDownChar(key)
            time.sleep(interval)
        for key in key_list[::-1]:
            self.keyUpChar(key)
            time.sleep(interval)

    # def keyPressStr(self, key_list, interval=None):
    #     """
    #     支持两种格式，一种是"a,b,c",一种是["a","b","c"],或("a","b","c")等可迭代对象
    #     :param key_list:
    #     :return:
    #     """
    #     if not interval:
    #         interval = self.mouse_delay
    #     if isinstance(key_list, str):
    #         for key in key_list.split(","):
    #             self.keyPressChar(key)
    #             time.sleep(interval)
    #     else:
    #         for key in key_list:
    #             self.keyPressChar(key)
    #             time.sleep(interval)

    # def enableRealMouse(self, enable, mousedelay, mousestep):
    #     pass

    @staticmethod
    def get_handle():
        """
            获取所有句柄
        :return:
        """
        hwnd_title = {}
        def get_all_hwnd(hwnd, mouse):
            if (win32gui.IsWindow(hwnd)
                    and win32gui.IsWindowEnabled(hwnd)
                    and win32gui.IsWindowVisible(hwnd)):
                hwnd_title.update({hwnd : win32gui.GetWindowText(hwnd)})

        win32gui.EnumWindows(get_all_hwnd, 0)
        return hwnd_title

    @staticmethod
    def get_child_handle(parent_hwnd):
        """
            获取当前句柄的子句柄
        :param parent_hwnd:
        :return:
        """
        hWndChildList = []  # 子弹框列表
        # 枚举所有子框
        win32gui.EnumChildWindows(parent_hwnd, lambda hWnd, param: param.append(hWnd), hWndChildList)
        hwnd_title = {}
        def get_all_hwnd(hwnd):
            if (win32gui.IsWindow(hwnd)
                    and win32gui.IsWindowEnabled(hwnd)
                    and win32gui.IsWindowVisible(hwnd)):
                hwnd_title.update({win32gui.GetWindowText(hwnd): hwnd})

        [get_all_hwnd(i) for i in hWndChildList]
        return hWndChildList, hwnd_title

    @staticmethod
    def windows_capture(hwnd):
        # 后台
        # pip install pywin32 -i https://pypi.douban.com/simple
        # 获取窗口宽，高
        x, y, x2, y2 = win32gui.GetWindowRect(hwnd)
        width = x2 - x
        height = y2 - y

        print(f"width:{width} height:{height} x:{x} y:{y} x2:{x2} y2:{y2}")

        # 截图
        hWndDC = win32gui.GetWindowDC(hwnd)
        mfcDC = win32ui.CreateDCFromHandle(hWndDC)
        saveDC = mfcDC.CreateCompatibleDC()
        saveBitMap = win32ui.CreateBitmap()
        saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
        saveDC.SelectObject(saveBitMap)
        saveDC.BitBlt((0, 0), (width, height), mfcDC, (0, 0), win32con.SRCCOPY)
        # 保存图片转cv图像
        signedIntsArray = saveBitMap.GetBitmapBits(True)
        im_opencv = np.frombuffer(signedIntsArray, dtype='uint8')
        im_opencv.shape = (height, width, 4)
        img = cv2.cvtColor(im_opencv, cv2.COLOR_BGRA2BGR)
        # 释放内存
        win32gui.DeleteObject(saveBitMap.GetHandle())
        saveDC.DeleteDC()
        mfcDC.DeleteDC()
        win32gui.ReleaseDC(hwnd, hWndDC)
        return img

    @staticmethod
    def windows_capture2(hwnd):
        """
            防止黑屏
        :param hwnd:
        :return:
        """
        windll.user32.SetProcessDPIAware()
        left, top, right, bottom = win32gui.GetClientRect(hwnd)
        w = right - left
        h = bottom - top
        hwnd_dc = win32gui.GetWindowDC(hwnd)
        mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
        save_dc = mfc_dc.CreateCompatibleDC()
        bitmap = win32ui.CreateBitmap()
        bitmap.CreateCompatibleBitmap(mfc_dc, w, h)
        save_dc.SelectObject(bitmap)
        # If Special K is running, this number is 3. If not, 1
        result = windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 3)
        bmpinfo = bitmap.GetInfo()
        bmpstr = bitmap.GetBitmapBits(True)
        img = np.frombuffer(bmpstr, dtype=np.uint8).reshape((bmpinfo["bmHeight"], bmpinfo["bmWidth"], 4))
        img = np.ascontiguousarray(img)[..., :-1]  # make image C_CONTIGUOUS and drop alpha channel
        if not result:  # result should be 1
            win32gui.DeleteObject(bitmap.GetHandle())
            save_dc.DeleteDC()
            mfc_dc.DeleteDC()
            win32gui.ReleaseDC(hwnd, hwnd_dc)
            raise RuntimeError(f"Unable to acquire screenshot! Result: {result}")
        #cv2.imwrite("xxx.png", img)
        return img

    @staticmethod
    def normal_capture(hwnd):
        # 前台
        if win32gui.GetWindowPlacement(hwnd)[1] != win32con.SW_SHOWNORMAL:
            win32gui.SendMessage(hwnd, win32con.WM_SYSCOMMAND, win32con.SC_RESTORE, 0)
        topHwnd = win32gui.GetForegroundWindow()
        if topHwnd != hwnd:
            try:
                win32gui.SetForegroundWindow(hwnd)
            except Exception as e:
                print(f"SetForegroundWindow error:{e}")
                raise "无法置顶窗口"  # 有时候无法置顶窗口

        # # 是否激活窗口,避免遮挡
        # if win32gui.GetWindowPlacement(self.hwnd)[1] != win32con.SW_SHOWNORMAL:
        #     win32gui.SendMessage(self.hwnd, win32con.WM_SYSCOMMAND, win32con.SC_RESTORE, 0)
        #     time.sleep(0.1)
        # # 是否显示窗口
        # topHwnd = win32gui.GetForegroundWindow()
        # if topHwnd != self.hwnd:
        #     win32gui.SetForegroundWindow(self.hwnd)
        #     time.sleep(0.3)
        # 截图
        desk_hwnd = win32gui.GetDesktopWindow()
        img = KM_windows.windows_capture(desk_hwnd)
        x1, y1, x2, y2 = win32gui.GetWindowRect(hwnd)
        img = img[y1:y2, x1:x2]
        return img

    @staticmethod
    def imshow(img, use_cv=False):
        if use_cv:
            cv2.imshow('result', img)
            cv2.waitKey(0)
        else:
            plt.imshow(img)
            plt.show()

    @staticmethod
    def capture(use_pil=False):
        beg = time.time()
        if use_pil:
            from PIL import ImageGrab
            # 截图
            img = ImageGrab.grab(bbox=None)
            img = np.array(img.getdata(), np.uint8).reshape(img.size[1], img.size[0], 3)
            end = time.time()
            cv2.imshow('result', img)
            cv2.waitKey(0)
        else:
            # 能截到扩展桌面
            # 获取桌面
            hwin = win32gui.GetDesktopWindow()
            width = win32api.GetSystemMetrics(win32con.SM_CXVIRTUALSCREEN)
            height = win32api.GetSystemMetrics(win32con.SM_CYVIRTUALSCREEN)
            left = win32api.GetSystemMetrics(win32con.SM_XVIRTUALSCREEN)
            top = win32api.GetSystemMetrics(win32con.SM_YVIRTUALSCREEN)
            print(f"width:{width} height:{height} left:{left} top:{top}")
            # 返回句柄窗口的设备环境，覆盖整个窗口，包括非客户区，标题栏，菜单，边框
            hwindc = win32gui.GetWindowDC(hwin)
            # 创建设备描述表
            srcdc = win32ui.CreateDCFromHandle(hwindc)
            # 创建内存设备描述表
            memdc = srcdc.CreateCompatibleDC()
            # 创建位图对象准备保存图片
            bmp = win32ui.CreateBitmap()
            # 为bitmap开辟存储空间
            bmp.CreateCompatibleBitmap(srcdc, width, height)
            # 将截图保存到saveBitMap中
            memdc.SelectObject(bmp)
            # 保存bitmap到内存设备描述表
            memdc.BitBlt((0, 0), (width, height), srcdc, (left, top), win32con.SRCCOPY)
            # 获取像数值
            signedIntsArray = bmp.GetBitmapBits(False)
            img = np.array(signedIntsArray).astype(dtype="uint8")  # This is REALLY slow!
            img.shape = (height, width, 4)
            end = time.time()
            # 内存释放
            srcdc.DeleteDC()
            memdc.DeleteDC()
            win32gui.ReleaseDC(hwin, hwindc)
            win32gui.DeleteObject(bmp.GetHandle())

            # 通过cv2 绘图
            img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
            cv2.imshow('result', img)
            cv2.waitKey(0)

        print(end - beg)


def test_txt():
    # 以一个记事本窗口为例子
    hwnd = win32gui.FindWindowEx(win32gui.FindWindow('Notepad', None), None, 'Edit', None)
    hWndChildList = []  # 子弹框列表
    # 枚举所有子框
    win32gui.EnumChildWindows(4460538, lambda hWnd, param: param.append(hWnd), hWndChildList)
    print(hWndChildList)
    count = 0
    for i in hWndChildList:  # 找到子框对应的索引
        print(win32gui.GetWindowText(i), count)
        # 遍历找到可以编辑的 句柄
        # win32api.SendMessage(i, win32con.WM_CHAR, ord(f"{count}"), 0)
        # count += 1
    win32gui.FindWindow('Notepad', None)
    win32api.SendMessage(4920344, win32con.WM_CHAR, ord("A"), 0)
    win32api.SendMessage(4920344, win32con.WM_CHAR, ord("我"), 0)


if __name__ == '__main__':

    handle = KM_windows.get_handle()
    [print(h, handle.get(h)) for h in handle]

    #img = KM_windows.capture()
    #KM_windows.imshow(img, True)

    # img = KM_windows.windows_capture(4460538)
    # KM_windows.imshow(img, False)
    #
    # img = KM_windows.windows_capture2(4460538)
    # KM_windows.imshow(img, False)

    test_txt()




